package com.treasure.infrastructure.domain.llm.qianwen.client;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.aigc.generation.models.QwenParam;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.dashscope.utils.Constants;
import com.treasure.common.model.SessionUser;
import com.treasure.common.model.SessionUserHolder;
import com.treasure.common.utils.TSON;
import com.treasure.common.model.ChatHistoryEntity;
import com.treasure.domain.domain.ChatHistoryDomainService;
import com.treasure.common.model.command.ConversationRequest;
import com.treasure.infrastructure.domain.llm.qianwen.DashscopeResultCallback;
import com.treasure.infrastructure.model.event.ConversationEvent;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.BooleanUtils;
import org.jetbrains.annotations.NotNull;
import org.springframework.context.ApplicationEventPublisher;
import org.springframework.context.ApplicationEventPublisherAware;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

/**
 * @author sunhongmin
 * @date 2023/8/23 19:56
 * @description
 */
@Service
@Slf4j
public class QwenGenerationClient extends QwenContext implements ApplicationEventPublisherAware {
    public static final String SSE_END_KEY = "SSE-END";
    @Resource
    private ChatHistoryDomainService chatHistoryDomainService;

    private ApplicationEventPublisher applicationEventPublisher;

    public GenerationResult qwenGeneration(String prompt) throws NoApiKeyException, InputRequiredException {
        log.info("qwenGeneration use model : " + qwenConfigBean.getModel());
        Constants.apiKey = qwenConfigBean.getKey();
        Generation gen = new Generation();
        Message systemMsg = Message.builder().role(Role.SYSTEM.getValue()).content("你是高校招生办老师").build();
        QwenParam param = QwenParam.builder()
                .messages(Collections.singletonList(systemMsg))
                .model(qwenConfigBean.getModel())
                .prompt(prompt)
                .topP(0.5)
                .build();
        return gen.call(param);
    }


    public String doGeneration(ConversationRequest conversationRequest) throws NoApiKeyException, InputRequiredException {
        SessionUser sessionUser = SessionUserHolder.getSessionUser();
        log.info("qwenGeneration use model : " + qwenConfigBean.getModel());
        Constants.apiKey = qwenConfigBean.getKey();
        // 组装聊天历史信息
        List<Message> historyMessages = historyMessageAssembler(conversationRequest.getChatHistoryList());
        // 请求模型
        log.info("qwenGenerationWithHistory historyMessage:{},prompt:{}", TSON.toJsonString(historyMessages), TSON.toJsonString(conversationRequest.getPrompt()));
        QwenParam param = QwenParam.builder()
                .messages(historyMessages)
                .model(qwenConfigBean.getModel())
                .prompt(conversationRequest.getPrompt().getValue())
                .resultFormat(QwenParam.ResultFormat.MESSAGE)
                .topP(0.3)// 生成过程中核采样方法概率阈值，例如，取值为0.8时，仅保留概率加起来大于等于0.8的最可能token的最小集合作为候选集。取值范围为（0,1.0)，取值越大，生成的随机性越高；取值越低，生成的确定性越高。
                .topK(20) // 生成时，采样候选集的大小。例如，取值为50时，仅将单次生成中得分最高的50个token组成随机采样的候选集。取值越大，生成的随机性越高；取值越小，生成的确定性越高。默认不传递该参数，取值为None或当top_k大于100时，表示不启用top_k策略，此时，仅有top_p策略生效。
                .seed(500) // 生成时使用的随机数种子，用户控制模型生成内容的随机性。seed支持无符号64位整数，默认值为1234。在使用seed时，模型将尽可能生成相同或相似的结果，但目前不保证每次生成的结果完全相同。
                .enableSearch(conversationRequest.getOpenSearch())
                .build();
        Generation generation = new Generation();
        if (BooleanUtils.isTrue(conversationRequest.getSyncCall())) {
            // 同步响应
            GenerationResult generationResult = generation.call(param);
            String response = generationResult.getOutput().getChoices().get(0).getMessage().getContent();
            log.info("回答： \n {}", response);
            // 记录聊天历史
            ChatHistoryEntity chatHistoryEntity = DashscopeResultCallback.getChatHistoryEntity(conversationRequest, response);
            chatHistoryDomainService.addChatHistory(chatHistoryEntity);
            int inputTokens = generationResult.getUsage().getInputTokens();
            int totalOutputTokens = generationResult.getUsage().getOutputTokens();
            ConversationEvent conversationEvent = DashscopeResultCallback.getConversationEvent(sessionUser, conversationRequest, new ConversationEvent(this), qwenConfigBean.getModel(), response, inputTokens, totalOutputTokens);
            applicationEventPublisher.publishEvent(conversationEvent);
            return response;
        }
        // 流式响应
        generation.streamCall(param, new DashscopeResultCallback(this, sessionUser, conversationRequest, chatHistoryDomainService, applicationEventPublisher));
        return "SUCCESS";
    }

    @NotNull
    private static List<Message> historyMessageAssembler(List<ChatHistoryEntity.ChatHistoryBean> historyBeanList) {
        // 历史聊天信息
        List<Message> historyMessages = new ArrayList<>(historyBeanList.size());

        // system 标识
        Message systemMsg = Message.builder().role(Role.SYSTEM.getValue()).content("你是高校招生办老师").build();
        historyMessages.add(systemMsg);
        if (CollectionUtils.isNotEmpty(historyBeanList)) {
            for (ChatHistoryEntity.ChatHistoryBean chatHistoryBean : historyBeanList) {
                historyMessages.add(Message.builder().role(chatHistoryBean.getIden()).content(chatHistoryBean.getMessage()).build());
            }
        }
        return historyMessages;
    }

    @Override
    public void setApplicationEventPublisher(@NotNull ApplicationEventPublisher applicationEventPublisher) {
        this.applicationEventPublisher = applicationEventPublisher;
    }
}
